{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建操作\n",
    "> 注意：为了保证你的 C++ 自定义操作系统与 tensorflow 的官方 PIP 软件包兼容，请遵循定制的 OP 存储库（[custom op repository](https://github.com/tensorflow/custom-op)）指南。它有一个端到端的代码示例，以及用于构建和分发自定义操作的 Docker 图像。\n",
    "\n",
    "如果你想创建一个现有的 tensorflow 库没有的 op，我们建议你首先尝试用 python 编写 op，作为现有python ops 或函数的组合。如果不可能，你可以创建一个自定义 C++ op。有几个原因可能需要创建一个自定义 C++ op：\n",
    "- 将你的操作表示为现有操作的组合是不容易或不可能的。\n",
    "- 将你的操作表示为现有的组合是没有效率的。\n",
    "- 你希望手工组装将来的编译器难以组装的原语组合。\n",
    "\n",
    "例如，假设你希望实现类似于 \"MaxPool\" 操作符的 \"median pooling\" 之类的东西，在计算滑动窗口的中位数而不是最大值。使用操作组合执行此操作可能是可能的（例如，使用 ExtractImagePatches 和TopK），但可能不如本机操作的性能或内存效率高，在本机操作中，你可以在单个融合操作中执行更灵活的操作。与往常一样，通常首先值得尝试使用运算符组合来表达你想要的内容，只有在证明这是困难或低效的情况下才选择添加新操作。\n",
    "\n",
    "要合并你的自定义op，你需要:\n",
    "1. 在 C++ 文件中注册新的 op。操作注册为操作的功能定义了一个接口（规范），它独立于操作的实现。例如，操作注册定义了操作的名称以及操作的输入和输出。它还定义了用于张量形状推断的形状函数。\n",
    "2. 在 C++ 中实现 op。实现 op 的称为内核，它是你在步骤 1 中注册的规范的具体实现。对于不同的输入/输出类型或体系结构（例如，CPU、GPU），可以有多个内核。\n",
    "3. 创建一个Python包装器（可选）。这个包装器是用于在 python 中创建 op 的公共 API。默认包装器是从操作注册生成的，可以直接使用或添加。\n",
    "4. 编写一个函数来计算 op 的梯度（可选）。\n",
    "5. 测试op。我们通常在 python 中为了方便而这样做，但是你也可以在 C++ 中进行测试。如果你定义了梯度，可以使用 python 的 tf.test.compute_gradient_error 验证它们。可参阅 [relu_op_test.py](https://tensorflow.google.cn/code/tensorflow/python/kernel_tests/relu_op_test.py) 作为测试类 relu 运算符的前向传播及其梯度的示例。\n",
    "\n",
    "## 预备知识\n",
    "- 熟悉 C++。\n",
    "- 必须已经安装了 tensorflow 二进制文件，或者已经下载了 tensorflow 源代码，并且能够构建它。\n",
    "\n",
    "## 定义 op 接口\n",
    "通过在 tensorflow 系统中注册 op 来定义它的接口。在注册中，你可以指定 op 的名称、它的输入（类型和名称）和输出（类型和名称），以及文档注释和 op 可能需要的任何属性。\n",
    "\n",
    "为了了解这是如何工作的，假设你想要创建一个 op ，它接受一个 int32 的张量，并输出一个张量的副本，除了第一个元素之外，其他元素都设置为零。为此，创建名为 zero_out.cc 的文件。然后添加对定义操作接口的 REGISTER-OP 宏的调用：\n",
    "```c++\n",
    "#include \"tensorflow/core/framework/op.h\"\n",
    "#include \"tensorflow/core/framework/shape_inference.h\"\n",
    "\n",
    "using namespace tensorflow;\n",
    "\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "    .Input(\"to_zero: int32\")\n",
    "    .Output(\"zeroed: int32\")\n",
    "    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {\n",
    "      c->set_output(0, c->input(0));\n",
    "      return Status::OK();\n",
    "    });\n",
    "```\n",
    "\n",
    "此 ZeroOut 以 32 位整数的一个张量到 0 作为输入，并输出 32 位整数的一个张量归零。op 还使用一个shape 函数来确保输出张量与输入张量是相同的形状。例如，如果输入是形状的张量 [10，20]，则此形状函数指定输出形状也是 [10，20]。\n",
    "\n",
    "> 注意：op 名称必须是骆驼命名法（CamelCase），并且在二进制文件中注册的所有其他 op 中它必须是唯一的。\n",
    "\n",
    "## 实现 op 的内核\n",
    "定义接口后，你需要提供操作的一个或多个实现。若要创建这些内核，应创建扩展操作内核并重写 Compute 方法的类。Compute 方法提供一个 OpKernelContext 类型的上下文参数，你可以从中访问有用的内容，如输入和输出张量。\n",
    "\n",
    "将你的内核添加到你在上面创建的文件中。内核可能是这样的：\n",
    "```c++\n",
    "#include \"tensorflow/core/framework/op_kernel.h\"\n",
    "\n",
    "using namespace tensorflow;\n",
    "\n",
    "class ZeroOutOp : public OpKernel {\n",
    " public:\n",
    "  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}\n",
    "\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // 获取输入张量\n",
    "    const Tensor& input_tensor = context->input(0);\n",
    "    auto input = input_tensor.flat<int32>();\n",
    "\n",
    "    // 创建输出张量\n",
    "    Tensor* output_tensor = NULL;\n",
    "    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),\n",
    "                                                     &output_tensor));\n",
    "    auto output_flat = output_tensor->flat<int32>();\n",
    "\n",
    "    // 把输出张量的第一个元素以外的所有元素都设为0\n",
    "    const int N = input.size();\n",
    "    for (int i = 1; i < N; i++) {\n",
    "      output_flat(i) = 0;\n",
    "    }\n",
    "\n",
    "    // 如果可能，保持第一个输入值\n",
    "    if (N > 0) output_flat(0) = input(0);\n",
    "  }\n",
    "};\n",
    "```\n",
    "\n",
    "在实现内核之后，将其注册到 tensorflow 系统。在注册中，你可以指定运行该内核的不同约束。例如，可以为 CPU 创建一个内核，为 GPU 创建一个内核。\n",
    "\n",
    "要对 ZeroOut 操作执行此操作，应将以下内容添加到 zero_out.cc:\n",
    "```c++\n",
    "REGISTER_KERNEL_BUILDER(Name(\"ZeroOut\").Device(DEVICE_CPU), ZeroOutOp);\n",
    "```\n",
    "\n",
    "> 重点：可以并发地访问 OpKernel 的实例。你的计算方法必须是线程安全的。用互斥锁保护对类成员的任何访问。或者更好不要通过类成员共享状态，考虑使用 [ResourceMgr](https://tensorflow.google.cn/code/tensorflow/core/framework/resource_mgr.h) 来跟踪 op 状态。\n",
    "\n",
    "### 多线程 CPU 内核\n",
    "要编写多线程 CPU 内核，可以使用 work_sharder.h 中的 Shard 函数。此函数在配置为用于操作内线程的线程之间共享一个计算函数（参阅 [config.proto](https://tensorflow.google.cn/code/tensorflow/core/protobuf/config.proto) 中的 intra_op_parallelism_threads）。\n",
    "\n",
    "### GPU 内核\n",
    "GPU 内核由两部分实现：OpKernel 和 CUDA 内核及其启动代码。\n",
    "\n",
    "有时，操作内核的实现在 CP U和 GPU 内核之间是通用的，例如检查输入和分配输出。在这种情况下，建议的实施方式是：\n",
    "1. 定义设备上的 OpKernel 模板和张量的基本类型。\n",
    "2. 要执行输出的实际计算，Compute 函数调用模板函数结构。\n",
    "3. CPUDevice 的那个函子的专门化定义在同一个文件中，但是 GPUDevice 的专门化定义在一个 .cu.cc 文件中，因为它将用 CUDA 编译器编译。\n",
    "\n",
    "下面是一个示例实现：\n",
    "```c++\n",
    "// kernel_example.h\n",
    "#ifndef KERNEL_EXAMPLE_H_\n",
    "#define KERNEL_EXAMPLE_H_\n",
    "\n",
    "template <typename Device, typename T>\n",
    "struct ExampleFunctor {\n",
    "  void operator()(const Device& d, int size, const T* in, T* out);\n",
    "};\n",
    "\n",
    "#if GOOGLE_CUDA\n",
    "// GpuDevice 的部分专用函子。\n",
    "template <typename Eigen::GpuDevice, typename T>\n",
    "struct ExampleFunctor {\n",
    "  void operator()(const Eigen::GpuDevice& d, int size, const T* in, T* out);\n",
    "};\n",
    "#endif\n",
    "\n",
    "#endif KERNEL_EXAMPLE_H_\n",
    "```\n",
    "```c++\n",
    "// kernel_example.cc\n",
    "#include \"kernel_example.h\"\n",
    "#include \"tensorflow/core/framework/op_kernel.h\"\n",
    "\n",
    "using namespace tensorflow;\n",
    "\n",
    "using CPUDevice = Eigen::ThreadPoolDevice;\n",
    "using GPUDevice = Eigen::GpuDevice;\n",
    "\n",
    "// 实际计算的 CPU 专门化。\n",
    "template <typename T>\n",
    "struct ExampleFunctor<CPUDevice, T> {\n",
    "  void operator()(const CPUDevice& d, int size, const T* in, T* out) {\n",
    "    for (int i = 0; i < size; ++i) {\n",
    "      out[i] = 2 * in[i];\n",
    "    }\n",
    "  }\n",
    "};\n",
    "\n",
    "// OpKernel 定义.\n",
    "// 模板参数<T>是张量的数据类型。\n",
    "template <typename Device, typename T>\n",
    "class ExampleOp : public OpKernel {\n",
    " public:\n",
    "  explicit ExampleOp(OpKernelConstruction* context) : OpKernel(context) {}\n",
    "\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // 获取输入张量\n",
    "    const Tensor& input_tensor = context->input(0);\n",
    "\n",
    "    // 创建一个输出张量\n",
    "    Tensor* output_tensor = NULL;\n",
    "    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),\n",
    "                                                     &output_tensor));\n",
    "\n",
    "    // 做计算\n",
    "    OP_REQUIRES(context, input_tensor.NumElements() <= tensorflow::kint32max,\n",
    "                errors::InvalidArgument(\"Too many elements in tensor\"));\n",
    "    ExampleFunctor<Device, T>()(\n",
    "        context->eigen_device<Device>(),\n",
    "        static_cast<int>(input_tensor.NumElements()),\n",
    "        input_tensor.flat<T>().data(),\n",
    "        output_tensor->flat<T>().data());\n",
    "  }\n",
    "};\n",
    "\n",
    "// 注册 CPU 内核。\n",
    "#define REGISTER_CPU(T)                                          \n",
    "  REGISTER_KERNEL_BUILDER(                                       \n",
    "      Name(\"Example\").Device(DEVICE_CPU).TypeConstraint<T>(\"T\"), \n",
    "      ExampleOp<CPUDevice, T>);\n",
    "REGISTER_CPU(float);\n",
    "REGISTER_CPU(int32);\n",
    "\n",
    "// 注册GPU内核\n",
    "#ifdef GOOGLE_CUDA\n",
    "#define REGISTER_GPU(T)                                          \n",
    "  /* 在 kernel_example.cu.cc 中声明显式实例化。 */ \n",
    "  extern template ExampleFunctor<GPUDevice, T>;                  \n",
    "  REGISTER_KERNEL_BUILDER(                                       \n",
    "      Name(\"Example\").Device(DEVICE_GPU).TypeConstraint<T>(\"T\"), \n",
    "      ExampleOp<GPUDevice, T>);\n",
    "REGISTER_GPU(float);\n",
    "REGISTER_GPU(int32);\n",
    "#endif  // GOOGLE_CUD\n",
    "```\n",
    "```c++\n",
    "// kernel_example.cu.cc\n",
    "#ifdef GOOGLE_CUDA\n",
    "#define EIGEN_USE_GPU\n",
    "#include \"example.h\"\n",
    "#include \"tensorflow/core/util/gpu_kernel_helper.h\"\n",
    "\n",
    "using namespace tensorflow;\n",
    "\n",
    "using GPUDevice = Eigen::GpuDevice;\n",
    "\n",
    "// 定义CUDA内核。\n",
    "template <typename T>\n",
    "__global__ void ExampleCudaKernel(const int size, const T* in, T* out) {\n",
    "  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;\n",
    "       i += blockDim.x * gridDim.x) {\n",
    "    out[i] = 2 * ldg(in + i);\n",
    "  }\n",
    "}\n",
    "\n",
    "// 定义启动 CUDA 内核的 GPU 实现。\n",
    "template <typename T>\n",
    "void ExampleFunctor<GPUDevice, T>::operator()(\n",
    "    const GPUDevice& d, int size, const T* in, T* out) {\n",
    "  // 启动 CUDA 内核。\n",
    "  //\n",
    "  // 有关计算块计数和 thread_per_block 计数的示例，参阅 core/util/gpu_kernel_helper.h。\n",
    "  int block_count = 1024;\n",
    "  int thread_per_block = 20;\n",
    "  ExampleCudaKernel<T>\n",
    "      <<<block_count, thread_per_block, 0, d.stream()>>>(size, in, out);\n",
    "}\n",
    "\n",
    "// 显式实例化已注册 opkerkernel 类型的函函数。\n",
    "template struct ExampleFunctor<GPUDevice, float>;\n",
    "template struct ExampleFunctor<GPUDevice, int32>;\n",
    "\n",
    "#endif  // GOOGLE_CUDA\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构建 op 库\n",
    "### 使用系统编译器编译 op  (tensorflow 二进制安装)\n",
    "你应该使用一个 C++ 编译器，例如 g++ 或 CLAN，编译 zero_out.cc。二进制 PIP 包在特定于系统的位置安装编译 op 所需的头文件和库。tensorflow python库提供 get_include 函数来获取头目录，get_lib目录有一个要链接的共享对象。以下是 Ubuntu 机器上这些函数的输出。\n",
    "```python\n",
    "import tensorflow as tf\n",
    "print(tf.sysconfig.get_include())\n",
    "print(tf.sysconfig.get_lib())\n",
    "```\n",
    "假设你已经安装了g++，下面是你可以用来将 op 编译成动态库的命令序列：\n",
    "```sh\n",
    "TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(\" \".join(tf.sysconfig.get_compile_flags()))') )\n",
    "TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(\" \".join(tf.sysconfig.get_link_flags()))') )\n",
    "g++ -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2\n",
    "```\n",
    "在 Mac OS X 上，在构建. so 文件时需要附加标记  \"-undefined dynamic_lookup\"。\n",
    "> 如果 gcc 版本 >= 5 请注意：gcc从版本5开始使用新的 C++ ABI。tensorflow 网站上提供的二进制 pip 包是使用 gcc4 构建的，该 gcc4 使用的是旧的 ABI。如果你使用 gcc>=5 编译op库，那么在命令行中添加 -D_GLIBCXX_USE_CXX11_ABI=0，使该库与旧的abi兼容。\n",
    "### 使用 bazel 编译 op（tensorflow 源文件安装）\n",
    "如果你安装了 tensorflow 源文件，你可以利用 tensorflow 的编译系统来编译你的操作。在 tensorflow /core/user_ops 目录下放置一个遵循 bazel 编译规则的编译文件。\n",
    "```python\n",
    "load(\"//tensorflow:tensorflow.bzl\", \"tf_custom_op_library\")\n",
    "\n",
    "tf_custom_op_library(\n",
    "    name = \"zero_out.so\",\n",
    "    srcs = [\"zero_out.cc\"],\n",
    ")\n",
    "```\n",
    "运行以下命令来构建zero_out.so：\n",
    "```sh\n",
    "bazel build --config opt //tensorflow/core/user_ops:zero_out.so\n",
    "```\n",
    "> 注意：如上所述，如果使用 gcc>=5 进行编译，则向Bazel命令行参数添加  --cxxopt=\"-D_GLIBCXX_USE_CXX11_ABI=0\"。\n",
    "\n",
    "> 注意：尽管可以使用标准的 cc_library 规则创建共享库（.so文件），但我们强烈建议你使用tf_custom_op_library 宏。它添加了一些必需的依赖项，并执行检查以确保共享库与 tensorflow的插件加载机制兼容。\n",
    "\n",
    "## 在 python 中使用 op\n",
    "tensorflow python API提供了 tf.load_op_library 函数，用于加载动态库并将 op 注册到 tensorflow 框架中。load_op_library 返回一个 python 模块，该模块包含用于 op 和内核的 python 包装器。因此，一旦你构建了op，你可以执行以下操作来从 python 运行它:\n",
    "```python\n",
    "import tensorflow as tf\n",
    "zero_out_module = tf.load_op_library('./zero_out.so')\n",
    "with tf.Session(''):\n",
    "  zero_out_module.zero_out([[1, 2], [3, 4]]).eval()\n",
    "\n",
    "# 输出\n",
    "array([[1, 0], [0, 0]], dtype=int32)\n",
    "```\n",
    "记住，生成的函数将被赋予 snake_case 名称（以符合 PEP8）。因此，如果你的 op 在 C++ 文件中被命名为 ZooOUT，python函数将被称为 zero_out.\n",
    "\n",
    "要使 op 作为可从 python 模块导入的常规函数可用，在 python 源文件中调用 load_op_library 可能会很有用，如下所示：\n",
    "```python\n",
    "import tensorflow as tf\n",
    "\n",
    "zero_out_module = tf.load_op_library('./zero_out.so')\n",
    "zero_out = zero_out_module.zero_out\n",
    "```\n",
    "## 验证 op 是否工作\n",
    "验证你已经成功实现了 op 的一个好方法是为它编写一个测试。创建包含以下内容的文件zero_out_op_test.py：\n",
    "```python\n",
    "import tensorflow as tf\n",
    "\n",
    "class ZeroOutTest(tf.test.TestCase):\n",
    "  def testZeroOut(self):\n",
    "    zero_out_module = tf.load_op_library('./zero_out.so')\n",
    "    with self.test_session():\n",
    "      result = zero_out_module.zero_out([5, 4, 3, 2, 1])\n",
    "      self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "  tf.test.main()\n",
    "```\n",
    "然后运行你的测试类：\n",
    "```sh\n",
    "python zero_out_op_test.py\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 在 op 中构建高级特性\n",
    "现在你已经了解了如何构建一个基本的（并且在一定程度上受到限制的）op 和实现，接下来我们将介绍一些你通常需要构建到 op 中的更复杂的东西：\n",
    "- 条件检查和验证\n",
    "- op 注册\n",
    "    - 属性\n",
    "    - 属性类型\n",
    "    - 多态性\n",
    "    - 输入和输出\n",
    "    - 向后兼容\n",
    "- 对GPU 的支持\n",
    "    - 为 GPU 设备编译内核\n",
    "- 在 python 中实现梯度变化\n",
    "- C++ 中的形状函数\n",
    "\n",
    "### 条件检查和验证\n",
    "上面的例子假设op应用于任意形状的张量。如果它只适用于向量呢？这意味着要向上面的 OpKernel 实现添加一个检查：\n",
    "```c++\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // 获取输入张量\n",
    "   const Tensor& input_tensor = context->input(0);\n",
    "\n",
    "    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),\n",
    "                errors::InvalidArgument(\"ZeroOut expects a 1-D vector.\"));\n",
    "    // ...\n",
    "  }\n",
    "```\n",
    "以上代码将断言输入是一个向量，如果不是，则返回设置了 InvalidArgument 状态的结果。[OP_REQUIRES](https://tensorflow.google.cn/code/tensorflow/core/lib/core/errors.h) 宏有三个参数:\n",
    "- context：上下文，它的 SetStatus() 方法可以是 OpKernelContext 或 OpKernelConstruction 指针(参见[tensorflow/core/framework/op_kernel.h](https://tensorflow.google.cn/code/tensorflow/core/framework/op_kernel.h))。\n",
    "- condition：条件，比如在 [tensorflow/core/framework/tensor_shape.h](https://tensorflow.google.cn/code/tensorflow/core/framework/tensor_shape.h) 中有一些函数可以验证张量的形状\n",
    "- error：错误本身由一个状态对象表示，参见 [tensorflow/core/lib/core/ Status.h](https://tensorflow.google.cn/code/tensorflow/core/lib/core/status.h)。状态既有类型（通常是InvalidArgument，但也可参阅类型列表），也有消息。构造错误的函数可以在 [tensorflow/core/lib/core/errors.h](https://tensorflow.google.cn/code/tensorflow/core/lib/core/errors.h) 中找到。\n",
    "或者，如果你想测试从某个函数返回的状态对象是否为错误，如果是，则返回它，使用 OP_REQUIRES_OK。这两个宏都在函数出错时返回。\n",
    "\n",
    "### op 注册\n",
    "#### 属性\n",
    "op 可以有属性值（attrs），其值是在将 op 添加到图中时设置的。它们用于配置 op，它们的值可以在内核实现中访问，也可以在 op 注册中的输入和输出类型中访问。如果可能的话，最好使用输入而不是 attr，因为输入更灵活。这是因为 attr 是常量，必须在图构造时定义。相反，输入是张量，其值可以是动态的；也就是说，输入可以改变每一步，使用 feed 设置，等等。attr 用于不能用输入完成的事情：任何影响签名（输入或输出的数量或类型）的配置，或者不能一步一步改变的配置。\n",
    "\n",
    "当你注册 op 时，通过使用 attr 方法指定它的名称和类型来定义 attr, attr 方法需要表单的规格：\n",
    "```txt\n",
    "<name>: <attr-type-expr>;\n",
    "```\n",
    "其中 &lt;name&gt; 以字母开头，可以由字母数字字符和下划线组成，&lt;attri-type-expr&gt; 是下面描述的表单的类型表达式。\n",
    "    \n",
    "例如，如果你想让 ZeroOut op 保存一个用户指定的索引，而不是只保存第 0 个元素，你可以这样注册 op：\n",
    "```c++\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "    .Attr(\"preserve_index: int\")\n",
    "    .Input(\"to_zero: int32\")\n",
    "    .Output(\"zeroed: int32\");\n",
    "```\n",
    "&gt; 注意：属性类型集（[the set of atttibute types](https://tensorflow.google.cn/guide/create_op#attr_types)）与用于输入和输出的 [tf.DType](https://tensorflow.google.cn/api_docs/python/tf/dtypes/DType) 不同。\n",
    "\n",
    "你的内核可以通过上下文参数在它的构造函数中访问这个 attr：\n",
    "```c++\n",
    "class ZeroOutOp : public OpKernel {\n",
    " public:\n",
    "  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {\n",
    "    // 获取要保存的值的索引\n",
    "    OP_REQUIRES_OK(context,\n",
    "                   context->GetAttr(\"preserve_index\", &preserve_index_));\n",
    "    // 检查p reserve_index 是否为正\n",
    "    OP_REQUIRES(context, preserve_index_ >= 0,\n",
    "                errors::InvalidArgument(\"Need preserve_index >= 0, got \",\n",
    "                                        preserve_index_));\n",
    "  }\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // ...\n",
    "  }\n",
    " private:\n",
    "  int preserve_index_;\n",
    "};\n",
    "```\n",
    "然后可以在 Compute 方法中使用：\n",
    "```c++\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // ...\n",
    "\n",
    "    // 我们使用已保存的 attr 来验证潜在的动态输入，以便检查 preserve_index 是否在范围内\n",
    "    OP_REQUIRES(context, preserve_index_  >  input.dimension(0),\n",
    "                errors::InvalidArgument(\"preserve_index out of range\"));\n",
    "\n",
    "    // 把输出张量的所有元素设为 0\n",
    "    const int N = input.size();\n",
    "    for (int i = 0; i > N; i++) {\n",
    "      output_flat(i) = 0;\n",
    "    }\n",
    "\n",
    "    // 保留请求的输入值\n",
    "    output_flat(preserve_index_) = input(preserve_index_);\n",
    "  }\n",
    "```\n",
    "\n",
    "#### 属性类型\n",
    "属性支持以下的类型：\n",
    "- string：任何字节序列（不要求是UTF8）。\n",
    "- int：一个单独的值。\n",
    "- float：一个浮点数\n",
    "- bool：true 或者 false。\n",
    "- type：[DataTtype](https://tensorflow.google.cn/code/tensorflow/core/framework/types.cc) 的非引用值之一\n",
    "- shape：一个 [TensorShapeProto](https://tensorflow.google.cn/code/tensorflow/core/framework/tensor_shape.proto)。\n",
    "- tensor：一个 [TensorProto](https://tensorflow.google.cn/code/tensorflow/core/framework/tensor.proto)。\n",
    "- list(&lt;type&gt;)：其中&lt;type&gt;是上述类型之一。注意 list(list(&lt;type&gt;)) 是无效的。\n",
    "\n",
    "可参见 [op_def_builder.cc:FinalizeAttr](https://tensorflow.google.cn/code/tensorflow/core/framework/op_def_builder.cc) 获取标准列表。\n",
    "\n",
    "**默认值和约束**：\n",
    "\n",
    "attrs 可能具有默认值，某些类型的 attrs 可能具有约束。要定义一个带有约束的 attr，你可以使用以下 &lt;attri -type-expr&lg;：\n",
    "- {'&lt;string1&gt;', '&lt;string2&gt;'}： 该值必须是一个值&lt;string1&gt;或&lt;string2&gt;的字符串。在使用这种语法时，类型的名称 string 是隐含的：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"EnumExample\")\n",
    "    .Attr(\"e: {'apple', 'orange'}\");\n",
    "    ```\n",
    "- {&lt;type1&gt;, &lt;type2&gt;}：该值的类型为type，必须是&lt;type1&gt;或&lt;type2&gt;中的一个，其中&lt;type1&gt;和&lt;type2&gt;都支持 tf.DType。你不需要指定 attr 的类型为 type。当你在 {…} 中有一个类型列表时，就会暗示这一点。例如，在这种情况下，attr t 是一种类型，它必须是 int32、float 或 bool：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"RestrictedTypeExample\")\n",
    "    .Attr(\"t: {int32, float, bool}\");\n",
    "    ```\n",
    "- 常见类型约束有快捷方式：\n",
    "    - numbertype：Type 类型限制为数字（非字符串和非 bool）类型。\n",
    "    - realnumbertype：没有复杂的类型，比如 numbertype。\n",
    "    - quantizedtype：量子化的数字类型。比如 numbertype。\n",
    "   这些类型所允许的特定列表是由 tensorflow/core/framework/types.h 中的函数（如 NumberTypes()）定义的。在本例中，attr t 必须是数字类型之一：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"NumberType\")\n",
    "    .Attr(\"t: numbertype\");\n",
    "    tf.number_type(t=tf.int32)  # 合法\n",
    "    tf.number_type(t=tf.bool)   # 不合法\n",
    "    ```\n",
    "    列表可以与其他列表和单一类型相结合。下面的操作允许 attr  t 是任何数字类型，或 bool 类型：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"NumberOrBooleanType\")\n",
    "        .Attr(\"t: {numbertype, bool}\");\n",
    "    tf.number_or_boolean_type(t=tf.int32)  # 合法\n",
    "    tf.number_or_boolean_type(t=tf.bool)   # 合法\n",
    "    tf.number_or_boolean_type(t=tf.string) # 不合法\n",
    "    ```\n",
    "    - int &gt;= &lt;n&gt;：该值必须是一个 int 值且大于或等于&lt;n&gt;，其中&lt;n&gt;是一个自然数。\n",
    "    例如，下面的 op 注册指定 attr a 必须有一个至少为 2 的值：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"MinIntExample\")\n",
    "        .Attr(\"a: int >= 2\");\n",
    "    ```\n",
    "    - 列表 (&lt;type&gt;) &gt;= &lt;n&gt;：长度大于或等于&lt;n&gt;的类型&lt;type&gt;的列表。\n",
    "    例如，下面的 op 注册指定 attr a 是类型列表（int32 或 float），并且必须至少有3种类型：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"TypeListExample\")\n",
    "        .Attr(\"a: list({int32, float}) >= 3\");\n",
    "    ```\n",
    "    要设置attr的默认值(使其在生成的代码中可选)，在最后 add = &lt;default&gt;，如：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"AttrDefaultExample\")\n",
    "        .Attr(\"i: int = 0\");\n",
    "    ```\n",
    "    此外，约束和默认值都可以指定：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"AttrConstraintAndDefaultExample\")\n",
    "        .Attr(\"i: int >= 1 = 1\");\n",
    "    ```\n",
    "    默认值支持的语法将用于最终的 GraphDe f定义的原型表示。以下是如何为所有类型指定默认值的示例：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"AttrDefaultExampleForAllTypes\")\n",
    "       .Attr(\"s: string = 'foo'\")\n",
    "       .Attr(\"i: int = 0\")\n",
    "       .Attr(\"f: float = 1.0\")\n",
    "       .Attr(\"b: bool = true\")\n",
    "       .Attr(\"ty: type = DT_INT32\")\n",
    "       .Attr(\"sh: shape = { dim { size: 1 } dim { size: 2 } }\")\n",
    "       .Attr(\"te: tensor = { dtype: DT_INT32 int_val: 5 }\")\n",
    "       .Attr(\"l_empty: list(int) = []\")\n",
    "       .Attr(\"l_int: list(int) = [2, 3, 5, 7]\");\n",
    "    ```\n",
    "    特别要注意 type 类型的值使用 tf.DType。\n",
    "#### 多态性\n",
    "**多态性类型**：\n",
    "对于可以采用不同类型作为输入或产生不同输出类型的操作，你可以在操作注册的输入或输出类型中指定 attr。通常，你会为每种支持的类型注册一个 OpKernel。\n",
    "\n",
    "例如，如果你希望除了 int32 之外，ZeroOut op 也能在 float 上工作，那么你的 op 注册可能如下所示：\n",
    "```c++\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "    .Attr(\"T: {float, int32}\")\n",
    "    .Input(\"to_zero: T\")\n",
    "    .Output(\"zeroed: T\");\n",
    "```\n",
    "你的 op 注册现在指定输入的类型必须是 float 或 int32，并且它的输出将是相同的类型，因为两者都有类型 T。\n",
    "\n",
    "**命名**：\n",
    "输入、输出和属性通常应指定下划线命名名称。一个例外是用作输入类型或输出类型的 attrs。当 op 被添加到图中时，可以推断出这些属性，因此不会出现在 op 的函数中。例如，ZeroOut 的最后一个定义将生成一个 python 函数，该函数如下所示：\n",
    "```python\n",
    "def zero_out(to_zero, name=None):\n",
    "  \"\"\"...\n",
    "  参数:\n",
    "    to_zero：一个张量，必须是以下类型：\n",
    "        float32、int32\n",
    "    name：op 的名字（可选）\n",
    "\n",
    "  返回:\n",
    "    一个和 `to_zero` 一样类型的张量\n",
    "  \"\"\"\n",
    "```\n",
    "如果 to_zero 传递了一个 int32 张量，那么 T 就会自动设置为 int32（实际上是 DT_INT32）。那些推断的 attrs 被转换为大写或者驼峰名称。\n",
    "\n",
    "将此与具有确定输出类型的 attr 类型的 op 进行比较：\n",
    "```c++\n",
    "REGISTER_OP(\"StringToNumber\")\n",
    "    .Input(\"string_tensor: string\")\n",
    "    .Output(\"output: out_type\")\n",
    "    .Attr(\"out_type: {float, int32} = DT_FLOAT\");\n",
    "    .Doc(R\"doc(\n",
    "Converts each string in the input Tensor to the specified numeric type.\n",
    ")doc\");\n",
    "```\n",
    "在这种情况下，用户必须指定输出类型，就像在生成的 python 中一样：\n",
    "```python\n",
    "def string_to_number(string_tensor, out_type=None, name=None):\n",
    "  \"\"\"将输入张量中的每个字符串转换为指定的数字类型.\n",
    "\n",
    "  参数:\n",
    "    string_tensor: string 类型的张量\n",
    "    out_type: 可选的 tf.DType 来自：`tf.float32，tf.int32`。\n",
    "      默认为  `tf.float32`.\n",
    "    name: 操作的名称（可选）\n",
    "\n",
    "  返回:\n",
    "    `out_type` 类型的张量\n",
    "  \"\"\"\n",
    "```\n",
    "\n",
    "**多态类型的例子**：\n",
    "```c++\n",
    "#include \"tensorflow/core/framework/op_kernel.h\"\n",
    "\n",
    "class ZeroOutInt32Op : public OpKernel {\n",
    "  // 如之前一样\n",
    "};\n",
    "\n",
    "class ZeroOutFloatOp : public OpKernel {\n",
    " public:\n",
    "  explicit ZeroOutFloatOp(OpKernelConstruction* context)\n",
    "      : OpKernel(context) {}\n",
    "\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // 获取输入张量\n",
    "    const Tensor& input_tensor = context->input(0);\n",
    "    auto input = input_tensor.flat<float>();\n",
    "\n",
    "    // 创建输出张量\n",
    "    Tensor* output = NULL;\n",
    "    OP_REQUIRES_OK(context,\n",
    "                   context->allocate_output(0, input_tensor.shape(), &output));\n",
    "    auto output_flat = output->template flat<float>();\n",
    "      \n",
    "    // 把输出张量的所有元素设为 0\n",
    "    const int N = input.size();\n",
    "    for (int i = 0; i < N; i++) {\n",
    "      output_flat(i) = 0;\n",
    "    }\n",
    "\n",
    "    // 保留第一个输入值\n",
    "    if (N > 0) output_flat(0) = input(0);\n",
    "  }\n",
    "};\n",
    "\n",
    "// 注意，TypeConstraint<int32&>(T) 表示attr \"T\"（在上面的操作注册中定义的）必须是 int32 才能使用此模板实例化。\n",
    "REGISTER_KERNEL_BUILDER(\n",
    "    Name(\"ZeroOut\")\n",
    "    .Device(DEVICE_CPU)\n",
    "    .TypeConstraint<int32&>(\"T\"),\n",
    "    ZeroOutInt32Op);\n",
    "REGISTER_KERNEL_BUILDER(\n",
    "    Name(\"ZeroOut\")\n",
    "    .Device(DEVICE_CPU)\n",
    "    .TypeConstraint<float>(\"T\"),\n",
    "    ZeroOutFloatOp);\n",
    "```\n",
    "为了保持向后兼容性（[backwards compatibility](https://tensorflow.google.cn/guide/create_op#backwards_compatibility)），你应该在向现有 op 添加 attr 时指定一个 [默认值](https://tensorflow.google.cn/guide/create_op#default_values_and_constraints)：\n",
    "```c++\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "  .Attr(\"T: {float, int32} = DT_INT32\")\n",
    "  .Input(\"to_zero: T\")\n",
    "  .Output(\"zeroed: T\")\n",
    "```\n",
    "假设你想添加更多类型，比如 double：\n",
    "```c++\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "    .Attr(\"T: {float, double, int32}\")\n",
    "    .Input(\"to_zero: T\")\n",
    "    .Output(\"zeroed: T\");\n",
    "```\n",
    "不像上面那样用冗余代码编写另一个 OpKernel，通常可以使用 C++ 模板。每个重载仍然有一个内核注册（REGISTER_KERNEL_BUILDER 调用）：\n",
    "```c++\n",
    "template <typename T>\n",
    "class ZeroOutOp : public OpKernel {\n",
    " public:\n",
    "  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}\n",
    "\n",
    "  void Compute(OpKernelContext* context) override {\n",
    "    // 获取输入张量\n",
    "    const Tensor& input_tensor = context->input(0);\n",
    "    auto input = input_tensor.flat<float>();\n",
    "\n",
    "    // 创建输出张量\n",
    "    Tensor* output = NULL;\n",
    "    OP_REQUIRES_OK(context,\n",
    "                   context->allocate_output(0, input_tensor.shape(), &output));\n",
    "    auto output_flat = output->template flat<float>();\n",
    "\n",
    "    // 把输出张量的所有元素设为 0\n",
    "    const int N = input.size();\n",
    "    for (int i = 0; i < N; i++) {\n",
    "      output_flat(i) = 0;\n",
    "    }\n",
    "\n",
    "    // 保留第一个输入值\n",
    "    if (N > 0) output_flat(0) = input(0);\n",
    "  }\n",
    "};\n",
    "\n",
    "// 注意，TypeConstraint<int32>(T) 表示attr \"T\"（在上面的操作注册中定义）必须是 int32，才能使用此模板实例化。\n",
    "REGISTER_KERNEL_BUILDER(\n",
    "    Name(\"ZeroOut\")\n",
    "    .Device(DEVICE_CPU)\n",
    "    .TypeConstraint<int32>(\"T\"),\n",
    "    ZeroOutInt32Op);\n",
    "REGISTER_KERNEL_BUILDER(\n",
    "    Name(\"ZeroOut\")\n",
    "    .Device(DEVICE_CPU)\n",
    "    .TypeConstraint<float>(\"T\"),\n",
    "    ZeroOutFloatOp);\n",
    "```\n",
    "如果有多个重载，可以将注册放到宏中：\n",
    "```c++\n",
    "#include \"tensorflow/core/framework/op_kernel.h\"\n",
    "\n",
    "#define REGISTER_KERNEL(type)                                       \\\n",
    "  REGISTER_KERNEL_BUILDER(                                          \\\n",
    "      Name(\"ZeroOut\").Device(DEVICE_CPU).TypeConstraint<type>(\"T\"), \\\n",
    "      ZeroOutOp<type>)\n",
    "\n",
    "REGISTER_KERNEL(int32);\n",
    "REGISTER_KERNEL(float);\n",
    "REGISTER_KERNEL(double);\n",
    "\n",
    "#undef REGISTER_KERNEL\n",
    "```\n",
    "根据你注册内核的类型列表，你可以使用 [tensorflow/core/framework/register_types.h](https://tensorflow.google.cn/code/tensorflow/core/framework/register_types.h)：\n",
    "```c++\n",
    "#include \"tensorflow/core/framework/op_kernel.h\"\n",
    "#include \"tensorflow/core/framework/register_types.h\"\n",
    "\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "    .Attr(\"T: realnumbertype\")\n",
    "    .Input(\"to_zero: T\")\n",
    "    .Output(\"zeroed: T\");\n",
    "\n",
    "template <typename T>\n",
    "class ZeroOutOp : public OpKernel { ... };\n",
    "\n",
    "#define REGISTER_KERNEL(type)                                       \\\n",
    "  REGISTER_KERNEL_BUILDER(                                          \\\n",
    "      Name(\"ZeroOut\").Device(DEVICE_CPU).TypeConstraint<type>(\"T\"), \\\n",
    "      ZeroOutOp<type>)\n",
    "\n",
    "TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);\n",
    "\n",
    "#undef REGISTER_KERNEL\n",
    "```\n",
    "\n",
    "**列出输入和输出**：\n",
    "除了能够接受或生成不同类型之外，op 还可以使用或生成数量不定的张量。\n",
    "\n",
    "在下一个示例中，attr T 保存类型列表，并用作输入和输出的类型。输入和输出是该类型张量的列表（输出中张量的数量和类型与输入相同，因为两者都具有类型T）。\n",
    "```c++\n",
    "REGISTER_OP(\"PolymorphicListExample\")\n",
    "    .Attr(\"T: list(type)\")\n",
    "    .Input(\"in: T\")\n",
    "    .Output(\"out: T\");\n",
    "```\n",
    "还可以对列表中可以指定的类型设置限制。在下一个例子中，输入是一个float 和 double 张量的列表。例如，op 接受输入类型（float、double、float），在这种情况下，输出类型也将是（float、double、float）。\n",
    "```c++\n",
    "REGISTER_OP(\"ListTypeRestrictionExample\")\n",
    "    .Attr(\"T: list({float, double})\")\n",
    "    .Input(\"in: T\")\n",
    "    .Output(\"out: T\");\n",
    "```\n",
    "如果你想要一个列表中的所有张量都是相同的类型，你可以这样做：\n",
    "```c++\n",
    "REGISTER_OP(\"IntListInputExample\")\n",
    "    .Attr(\"N: int\")\n",
    "    .Input(\"in: N * int32\")\n",
    "    .Output(\"out: int32\");\n",
    "```\n",
    "\n",
    "它接受一个 int32 张量的列表，并使用一个 int attr N 来指定列表的长度。\n",
    "\n",
    "这也可以使类型多态。在下一个例子中，输入是一个相同（但未指定）类型(\"T\")的张量列表（长度为“N”），输出是一个匹配类型的张量：\n",
    "```c++\n",
    "REGISTER_OP(\"SameListInputExample\")\n",
    "    .Attr(\"N: int\")\n",
    "    .Attr(\"T: type\")\n",
    "    .Input(\"in: N * T\")\n",
    "    .Output(\"out: T\");\n",
    "```\n",
    "默认情况下，张量列表的最小长度是 1。莫可以使用对应 attr 上的 \"&gt;=\" 约束来更改该缺省值。在下一个例子中，输入是一个至少包含 2 个 int32 张量的列表：\n",
    "```c++\n",
    "REGISTER_OP(\"MinLengthIntListExample\")\n",
    "    .Attr(\"N: int >= 2\")\n",
    "    .Input(\"in: N * int32\")\n",
    "    .Output(\"out: int32\");\n",
    "```\n",
    "\"list(type)\" attrs 也有相同的语法：\n",
    "```c++\n",
    "REGISTER_OP(\"MinimumLengthPolymorphicListExample\")\n",
    "    .Attr(\"T: list(type) >= 3\")\n",
    "    .Input(\"in: T\")\n",
    "    .Output(\"out: T\");\n",
    "```\n",
    "**输入和输出**：\n",
    "综上所述，一个 op 注册可以有多个输入和输出：\n",
    "```c++\n",
    "REGISTER_OP(\"MultipleInsAndOuts\")\n",
    "    .Input(\"y: int32\")\n",
    "    .Input(\"z: float\")\n",
    "    .Output(\"a: string\")\n",
    "    .Output(\"b: int32\");\n",
    "```\n",
    "每个输入或输出规格的形式如下：\n",
    "```txt\n",
    "<name>:<io-type-expr>\n",
    "```\n",
    "其中&lt;name&gt;以字母开头，可以由字母数字字符和下划线组成。&lt;io-type-expr&gt;是以下类型表达式之一：\n",
    "- &lt;type&gt;，其中&lt;type&gt;是受支持的输入类型（例如 float、int32、string）。它指定了给定类型的一个张量。\n",
    "    查看 [tf.DType](https://tensorflow.google.cn/api_docs/python/tf/dtypes/DType)：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"BuiltInTypesExample\")\n",
    "        .Input(\"integers: int32\")\n",
    "        .Input(\"complex_numbers: complex64\");\n",
    "    ```\n",
    "- &lt;attrtype&gt;，其中&lt;attrtype&gt;是具有类型类型或列表（类型）的 Attr 的名称（可能有类型限制）。这种语法允许多态操作。\n",
    "    ```c++\n",
    "    REGISTER_OP(\"PolymorphicSingleInput\")\n",
    "        .Attr(\"T: type\")\n",
    "        .Input(\"in: T\");\n",
    "\n",
    "    REGISTER_OP(\"RestrictedPolymorphicSingleInput\")\n",
    "        .Attr(\"T: {int32, int64}\")\n",
    "        .Input(\"in: T\");\n",
    "    ```\n",
    "    引用类型列表（type）的 attr 允许你接受一个张量序列：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"ArbitraryTensorSequenceExample\")\n",
    "        .Attr(\"T: list(type)\")\n",
    "        .Input(\"in: T\")\n",
    "        .Output(\"out: T\");\n",
    "\n",
    "    REGISTER_OP(\"RestrictedTensorSequenceExample\")\n",
    "        .Attr(\"T: list({int32, int64})\")\n",
    "        .Input(\"in: T\")\n",
    "        .Output(\"out: T\");\n",
    "    ```\n",
    "    注意，输出输出的张量的数量和类型与输入输入的张量相同，因为它们都属于 T 类型。\n",
    "    - 对于具有相同类型的张量序列：&lt;number&gt; * &lt;type&gt;，其中&lt;number&gt;是类型为 int 的 Attr 的名称。&lt;type&gt;可以是tf.DType，或具有类型类型的attr的名称。作为第一个例子，这个 op 接受一个 int32 张量列表：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"Int32SequenceExample\")\n",
    "        .Attr(\"NumTensors: int\")\n",
    "        .Input(\"in: NumTensors * int32\")\n",
    "    ```\n",
    "    而这个 o p接受任何类型的张量列表，只要它们都是相同：\n",
    "    ```c++\n",
    "    REGISTER_OP(\"SameTypeSequenceExample\")\n",
    "        .Attr(\"NumTensors: int\")\n",
    "        .Attr(\"T: type\")\n",
    "        .Input(\"in: NumTensors * T\")\n",
    "    ```\n",
    "    - 引用一个张量：Ref(&lt;type&gt;)，其中&lt;type&gt;是前面的类型之一。\n",
    "在输入类型中使用的任何 attr 都将被推断出来。按照惯例，推断的 attrs 使用大写名（如 T 或 N），否则输入、输出和 attrs 的名称类似于函数参数(如num_output)。有关更多细节，请参阅前面关于命名的部分。\n",
    "\n",
    "很多信息可查看 [tensorflow/core/framework/op_def_builder.h](https://tensorflow.google.cn/code/tensorflow/core/framework/op_def_builder.h)。\n",
    "\n",
    "**向后兼容**：\n",
    "让我们假设你已经编写了一个良好的自定义 op 并与他人共享它，这样你就有了满意的客户在使用你的操作。但是，你希望以某种方式对 op 进行更改。\n",
    "\n",
    "一般来说，对现有的、签入的规范的更改必须向后兼容：更改 op 的规范不能破坏以前的序列化 GraphDef 协议缓冲区，这些缓冲区是由以前的规范构造的。下面描述了 GraphDef 兼容性的细节：\n",
    "1. 添加到操作中的任何新 attrs 都必须定义默认值，使用该默认值，op 必须具有原始行为。要将一个操作从非多态更改为多态，必须为新类型 attr 提供一个默认值，以便在默认情况下保留原始签名。例如，如果你的操作是：\n",
    "```c++\n",
    "REGISTER_OP(\"MyGeneralUnaryOp\")\n",
    "    .Input(\"in: float\")\n",
    "    .Output(\"out: float\");\n",
    "# 你可以使它多态在一个向后兼容的方式使用:\n",
    "REGISTER_OP(\"MyGeneralUnaryOp\")\n",
    "    .Input(\"in: T\")\n",
    "    .Output(\"out: T\")\n",
    "    .Attr(\"T: numerictype = DT_FLOAT\");\n",
    "```\n",
    "2. 你可以安全地减少对 attr 的约束。例如，你可以将 {int32, int64} 更改为 {int32, int64, float} 或 type。或者你可以将 {\"apple\",\"orange\"} 更改为 {\"apple\",\"banana\",\"orange\"} 或字符串。\n",
    "3. 你可以将单个输入/输出更改为列表输入/输出，只要列表类型的默认值与旧签名匹配即可。\n",
    "4. 如果默认为空，你可以添加一个新的列表输入/输出。\n",
    "5. 命名你创建的任何新操作，方法是在操作名称前面加上项目特有的前缀。这避免了你的 op 与 tensorflow 未来版本中可能包含的任何 op 发生冲突。\n",
    "6. 提前计划，尝试预测该操作的未来用途。一些签名更改无法以兼容的方式完成（例如，将相同类型的列表转换为不同类型的列表）。\n",
    "\n",
    "安全与不安全更改的完整列表可以在 [tensorflow/core/framework/op_compatibility_test.cc](https://tensorflow.google.cn/code/tensorflow/core/framework/op_compatibility_test.cc) 中找到。如果你无法使对操作的更改向后兼容，则使用新语义的新名称创建新操作。\n",
    "\n",
    "还要注意，虽然这些更改可以保持 GraphDef 兼容性，但生成的 python 代码可能会以与旧调用者不兼容的方式更改。python API 可以通过在手工编写的 python 包装器中进行仔细的更改来保持兼容，除了可能在末尾添加新的可选参数外，还可以保留旧的签名。通常不兼容的更改只能在 tensorflow 更改主要版本时进行，并且必须符合 GraphDef 版本语义（[GraphDef version semantics](https://tensorflow.google.cn/guide/versions#compatibility_of_graphs_and_checkpoints)）。\n",
    "\n",
    "### 对 GPU 的支持\n",
    "你可以实现不同的 Opkerkernel，为 CPU 注册一个，为 GPU 注册另一个，就像你可以注册不同类型的内核一样。在 [tensorflow/core/kernels/](ensorflow/core/kernels/) 中有几个使用GPU支持内核的例子。注意，有些内核在 .cc 文件中有一个 CPU 版本，在  \\_gpu.cu 文件中有一个 GPU 版本。cc，以及 .h 文件中共享的一些代码。\n",
    "\n",
    "例如，[tf.pad](https://tensorflow.google.cn/api_docs/python/tf/pad) 除了 [tensorflow/core/kernels/pad_op.cc](https://tensorflow.google.cn/code/tensorflow/core/kernels/pad_op.cc) 中的 GPU 内核之外，还有其他东西。GPU 内核位于 [tensorflow/core/kernels/pad-op-GPU.cu.cc](https://tensorflow.google.cn/code/tensorflow/core/kernels/pad_op_gpu.cu.cc) 中，共享代码是[tensorflow/core/kernels/pad-op.h](https://tensorflow.google.cn/code/tensorflow/core/kernels/pad_op.h) 中定义的一个模板类。我们这样组织代码有两个原因：它允许你在 CPU 和 GPU 实现之间共享公共代码，它将 GPU 实现放在一个单独的文件中，这样它只能由 GPU 编译器编译。\n",
    "\n",
    "需要注意的是，即使使用的是 GP U内核版本的 pad，它仍然需要在 CPU 内存中输入 \"paddings\"。为了标记输入或输出保持在 CPU 上，添加一个 HostMemory() 调用到内核注册，例如：\n",
    "```c++\n",
    "#define REGISTER_GPU_KERNEL(T)                         \n",
    "  REGISTER_KERNEL_BUILDER(Name(\"Pad\")                  \n",
    "                              .Device(DEVICE_GPU)      \n",
    "                              .TypeConstraint<T>(\"T\")  \n",
    "                              .HostMemory(\"paddings\"), \n",
    "                          PadOp<GPUDevice, T>)\n",
    "```\n",
    "#### 为 GPU 设备编译内核\n",
    "查看 [cuda_op_kernel.cu.cc](https://tensorflow.google.cn/code/tensorflow/examples/adding_an_op/cuda_op_kernel.cu.cc) 以获取一个使用 CUDA 内核实现 op 的示例tf_custom_op_library 接受一个 gpu_srcs 参数，其中可以指定包含 CUDA 内核的源文件列表（.cu.cc文件）。要与 tensorflow 的二进制安装一起使用，必须使用 NVIDIA的 nvcc 编译器编译 CUDA 内核。下面是用于将 [cuda_op_kernel.cu.cc](https://tensorflow.google.cn/code/tensorflow/examples/adding_an_op/cuda_op_kernel.cu.cc) 和 [cuda_op_kernel.cc](https://tensorflow.google.cn/code/tensorflow/examples/adding_an_op/cuda_op_kernel.cc) 编译为单个动态可加载库的命令序列：\n",
    "```txt\n",
    "nvcc -std=c++11 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \\\n",
    "  ${TF_CFLAGS[@]} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC\n",
    "\n",
    "g++ -std=c++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \\\n",
    "  cuda_op_kernel.cu.o \\${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面生成的 cuda_op_kernel.so 可以像往常一样在 python 中使用 tf.load_op_library 函数加载。\n",
    "\n",
    "> 注意：如果CUDA库没有安装在/usr/local/lib64中，则需要在上面的第二个（g++）命令中显式地指定路径。例如，如果 CUDA 安装在 /usr/local/cuda-8.0 中，则添加 -L/usr/local/cuda-8.0/lib64/。\n",
    "\n",
    "> 注意：在某些 Linux 设置中，需要 nvcc 编译步骤的其他选项。在 nvcc 命令行中添加 -D_MWAITXINTRIN_H_INCLUDED，以避免来自 mwaitxintrin.h 的错误。\n",
    "\n",
    "## 在 python 中实现梯度变化\n",
    "\n",
    "给定一个 ops 图，tensorflow 使用自动微分（backpropagation）来添加表示相对于现有 ops 的梯度的新ops。要使新操作的自动区分工作，必须注册一个梯度函数，该函数计算相对于操作输入的梯度，给定相对于操作输出的梯度。\n",
    "\n",
    "在数学上，如果 op 计算出已注册的梯度，op 将通过链式法则将损失的梯度转化为梯度：\n",
    "\n",
    "在 ZeroOut 的情况下，输入中只有一个条目影响输出，因此相对于输入的梯度是稀疏的 \"一个热点\" 张量。具体表现如下：\n",
    "```python\n",
    "from tensorflow.python.framework import ops\n",
    "from tensorflow.python.ops import array_ops\n",
    "from tensorflow.python.ops import sparse_ops\n",
    "\n",
    "@ops.RegisterGradient(\"ZeroOut\")\n",
    "def _zero_out_grad(op, grad):\n",
    "  \"\"\"`zero_out` 的梯度\n",
    "\n",
    "  参数:\n",
    "    op: 我们正在微分的 zero_out 操作，我们可以使用它来查找原始 op 的输入和输出。\n",
    "    grad: 相对于 \"zero_out\" 操作的输出的梯度。\n",
    "\n",
    "  返回:\n",
    "   关于 zero_out输入的梯度。\n",
    "  \"\"\"\n",
    "  to_zero = op.inputs[0]\n",
    "  shape = array_ops.shape(to_zero)\n",
    "  index = array_ops.zeros_like(shape)\n",
    "  first_grad = array_ops.reshape(grad, [-1])[0]\n",
    "  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)\n",
    "  return [to_zero_grad]  # 一个张量的列表，因为我们有一个输入\n",
    "```\n",
    "\n",
    "关于使用 [tf.RegisterGradient](https://tensorflow.google.cn/api_docs/python/tf/RegisterGradient) 注册梯度函数的详细信息：\n",
    "1. 对于只有一个输出的 op，梯度函数将使用 [tf.Operation](https://tensorflow.google.cn/api_docs/python/tf/Operation) 的 op 和 [tf.Tensor](https://tensorflow.google.cn/api_docs/python/tf/Tensor) 的 grad，并根据 Tensor op.inputs[i]、op.outputs[i] 和 grad 构建新的 op。有关任何属性的信息可以通过 [tf.Operation.get_attr](https://tensorflow.google.cn/api_docs/python/tf/Operation#get_attr) 找到。\n",
    "2. 如果op有多个输出，梯度函数将采用 op 和 grads，其中 grads 是关于每个输出的梯度列表。梯度函数的结果必须是表示每个输入的梯度的张量对象列表。\n",
    "3. 如果某些输入没有定义良好的梯度，例如用作索引的整数输入，则相应的返回梯度应为 None。例如，对于采用浮点张量 x 和整数索引 i 的op，梯度函数将返回 [x_grad，None]。\n",
    "4. 如果op根本没有有意义的梯度，你通常不需要注册任何渐变，只要 op 的梯度不需要，就可以了。在某些情况下，op 没有明确定义的梯度，但可以参与梯度的计算。在这里，你可以使用 ops.NotDifferentiable 自动向后传播零。\n",
    "\n",
    "> 注意：在调用梯度函数时，只有 op 的数据流图可用，而张量数据本身不可用。因此，所有计算都必须使用其它 tensorflow 操作来执行，以在图形执行时运行。\n",
    "\n",
    "### C++ 中的形状函数\n",
    "TensorFlow API有一个称为 \"形状推断(shape inference)\" 的特性，它提供有关张量形状的信息，而无需执行图形。形状推理由在 C++ ReistSypOP 声明中注册的每个 op 类型支持的 \"形状函数\"，并执行两个角色：断言输入的形状在图形构造期间是兼容的，并且指定输出的形状。\n",
    "\n",
    "Shape 函数被定义为 shape_inference::InferenceContext 类上的操作。例如，在 shape function For ZeroOut 中：\n",
    "```c++\n",
    "    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {\n",
    "      c->set_output(0, c->input(0));\n",
    "      return Status::OK();\n",
    "    });\n",
    "```\n",
    "c->set_output(0, c->input(0))；声明第一个输出的形状应设置为第一个输入的形状。如果按照上面的示例中的索引选择输出，则 set_output 的第二个参数应该是 ShapeHandle 对象。可以使用默认构造函数创建空 ShapeHandle 对象。索引为 idx 的输入的 ShapeHandle 对象可以通过 c->input(dx) 获得。\n",
    "\n",
    "有许多常见的形状函数适用于许多操作，例如 shape_inference::UnchangedShape，可以在[common_shape_fns.h](https://tensorflow.google.cn/code/tensorflow/core/framework/common_shape_fns.h) 中找到，并使用如下：\n",
    "```c++\n",
    "REGISTER_OP(\"ZeroOut\")\n",
    "    .Input(\"to_zero: int32\")\n",
    "    .Output(\"zeroed: int32\")\n",
    "    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);\n",
    "```\n",
    "形状函数还可以约束输入的形状。对于具有矢量形状约束的归零版本（[ZeroOut with a vector shape constraint](https://tensorflow.google.cn/guide/create_op#conditional_checks_and_validation)），形状函数如下：\n",
    "```c++\n",
    "    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {\n",
    "      ::tensorflow::shape_inference::ShapeHandle input;\n",
    "      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input));\n",
    "      c->set_output(0, input);\n",
    "      return Status::OK();\n",
    "    });\n",
    "```\n",
    "WithRank 调用验证输入形状 c->input(0) 是否具有恰好一维的形状（或者如果输入形状未知，则输出形状将是一维未知的向量）。\n",
    "\n",
    "如果你的 op 具有多个输入的多态性，则可以使用 InferenceContext 的成员来确定要检查的形状数，并使用 Merge 来验证这些形状是否都兼容（或者，使用 InferenceContext::GetAttr 访问表示长度的属性，后者提供对 op 属性的访问）。\n",
    "```c++\n",
    "    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {\n",
    "      ::tensorflow::shape_inference::ShapeHandle input;\n",
    "      ::tensorflow::shape_inference::ShapeHandle output;\n",
    "      for (size_t i = 0; i < c->num_inputs(); ++i) {\n",
    "        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 2, &input));\n",
    "        TF_RETURN_IF_ERROR(c->Merge(output, input, &output));\n",
    "      }\n",
    "      c->set_output(0, output);\n",
    "      return Status::OK();\n",
    "    });\n",
    "```\n",
    "由于形状推断是一个可选的特征，张量的形状可以动态变化，因此形状函数必须对任何输入的不完整形状信息具有鲁棒性。InferenceContext 中的 Merge 方法允许调用者断言两个形状是相同的，即使其中一个或两个形状都没有完整的信息。定义了所有核心 tensorflow 操作的形状函数，并提供了许多不同的使用示例。\n",
    "\n",
    "InferenceContext 类有许多函数可用于定义形状函数操作。例如，可以使用 InferenceContext:: Dim和InferenceContext::WithValue 验证特定维度是否具有非常特定的值；可以使用 InferenceContext::Add 和 InferenceContext::Multiply 指定输出维度是两个输入维度的和/积。有关可以指定的所有形状操作，参见InferenceContext类。下面的示例将第一个输出的形状设置为（n，3），其中第一个输入具有形状（n，…）\n",
    "```c++\n",
    ".SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {\n",
    "    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));\n",
    "    return Status::OK();\n",
    "});\n",
    "```\n",
    "如果有复杂的形状函数，则应考虑添加一个测试，以验证各种输入形状组合是否生成预期的输出形状组合。你可以在一些核心操作测试（[core ops test](https://tensorflow.google.cn/code/tensorflow/core/ops/array_ops_test.cc)）中看到如何编写这些测试的示例。（INFER OK 和 inferu ERROR的语法 有点晦涩，但是在测试中尽量简洁地表示输入和输出形状规范。现在，请参阅这些测试中的周围注释以了解形状字符串规范）。\n",
    "\n",
    "## 为你的自定义 op 构建一个 pip 包\n",
    "要为 op 构建 pip 包，请参阅 [tensorflow/custom op](https://github.com/tensorflow/custom-op) 示例。本指南介绍如何 tTensorflow pip 包构建自定义操作，而不是从源代码构建 tensorflow。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
